# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Backbone modules.
"""
from collections import OrderedDict

import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from typing import Dict, List

from util.misc import NestedTensor, is_main_process

from .position_encoding import build_position_encoding


class FrozenBatchNorm2d(torch.nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters are fixed.

    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
    without which any other models than torchvision.models.resnet[18,34,50,101]
    produce nans.
    """

    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(n))
        self.register_buffer("bias", torch.zeros(n))
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]

        super(FrozenBatchNorm2d, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

    def forward(self, x):
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()
        bias = b - rm * scale
        return x * scale + bias


class BackboneBase(nn.Module):

    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
        super().__init__()
        for name, parameter in backbone.named_parameters():  # (python中优先级 not>or>and)
            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)  # 如果train_backbone为false，或者layer2,3,4都不在backbone中，backbone是用字典表示的，把backbone冻结，不进行梯度回传训练
        if return_interm_layers:  # 如果则返回中间层特征
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)  # IntermediateLayerGetter(Model)获取一个Model中你指定要获取的哪些层的输出，然后这些层的输出会在一个有序的字典中
        self.num_channels = num_channels

    def forward(self, tensor_list: NestedTensor):  #  tensor_list 是一个由张量和掩码组成的输入元组，即 NestedTensor(tensor, mask)  # tensor_list.tensors.Size([2, 3, 768, 1151])
        xs = self.body(tensor_list.tensors)  # 作用是将输入数据传递给模型的主体部分，进行特征提取或者其他计算操作，并将计算结果保存在 xs 中以供后续使用，即输入的tensor list 经过backbone后得到featuremap #  --> xs.Size([2, 2048, 24, 36])
        out: Dict[str, NestedTensor] = {}  # 作用是创建一个空的字典，用于存储模型输出的特征图或其他相关信息，并显式地指定其值的类型为 键值[str, NestedTensor]的字典
        for name, x in xs.items():  # x是其中一层的特征张量，name是该层的名称   # name:'0', x.Size([2, 2048, 24, 36])
            m = tensor_list.mask  # 获取输入 tensor_list 中的掩码信息，用于指示每个特征图中有效像素的位置  # m.Size([2, 768, 1151])
            assert m is not None
            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]  # 将输入的掩码信息进行插值操作，使其与当前特征图大小相匹配，并将结果转换为布尔类型的张量，以表示有效像素的位置。x.shape[-2:]表示在最后两个维度(高宽)进行  # mask.Size([2, 24, 36])
            out[name] = NestedTensor(x, mask)  # 将图像张量与mask封装到一起
        return out  # 返回(x, mask)，（mask 用于数据增强？× 把所有图片统一大小√）


class Backbone(BackboneBase):  # 构建骨干网络，父类为BackboneBase
    """ResNet backbone with frozen BatchNorm."""
    def __init__(self, name: str,
                 train_backbone: bool,  # 是否训练backbone
                 return_interm_layers: bool,  # 是否返回中间层特征
                 dilation: bool):  # 是否使用空洞卷积
        # getattr（obj,name）获取obj中命名为name的组成。可以理解为获取obj.name
        backbone = getattr(torchvision.models, name)(  # torchvision.models是pytorch的一个重要的包，包含了各种常用的网络结构，并且提供了预训练模型
            replace_stride_with_dilation=[False, False, dilation],
            weights="IMAGENET1K_V1",  # # torchvision 的 pretrained 参数已经弃用，从 0.13 版本开始，建议使用 weights 参数代替
            # pretrained=is_main_process(), # torchvision 的 pretrained 参数已经弃用，从 0.13 版本开始，建议使用 weights 参数代替
            norm_layer=FrozenBatchNorm2d)  # FrozenBatchNorm2d是一个冻结的BatchNorm层，其参数不会被更新，适用于迁移学习中固定特征提取器的情况
        num_channels = 512 if name in ('resnet18', 'resnet34') else 2048  # 对于resnet18和34，输出通道数为512，而对于其他类型（如resnet50、101等），输出通道数为2048
        super().__init__(backbone, train_backbone, num_channels, return_interm_layers)  # 调用父类BackboneBase的构造函数，并传递参数


class Joiner(nn.Sequential):  # 将 backbone 和 position_embedding 结合并读取
    def __init__(self, backbone, position_embedding):
        super().__init__(backbone, position_embedding)

    def forward(self, tensor_list: NestedTensor):
        xs = self[0](tensor_list)  # self[0] 指的是 Joiner 类的第一个子模块，即 backbone
        out: List[NestedTensor] = []
        pos = []
        for name, x in xs.items():  # x 是其中一层的特征张量，name 是该层的名称
            out.append(x)
            # position encoding
            pos.append(self[1](x).to(x.tensors.dtype))  # self[1](x) 调用了 position_embedding 子模块，即对特征张量 x 进行位置编码；to(x.tensors.dtype) 将位置编码张量转换为与输入特征张量 x 相同的数据类型

        return out, pos  # 返回list格式，包含多个需要返回的层数 # out[mask.Size([2, 24, 36]), tensor.Size([2, 2048, 24, 36])]，pos[torch.Size([2, 256, 24, 36])]

# ----------------阅读完毕--------------------先阅读build_model构建模型部分---->阅读build_backbone--->阅读Backbone-->阅读Joiner
def build_backbone(args):  # 搭建 backbone 网络
    position_embedding = build_position_encoding(args)  # 构建位置编码---未细看
    train_backbone = args.lr_backbone > 0  # 布尔值，传入Backbone()
    return_interm_layers = args.masks  # 是否要记录backbone每层的输出，传入Backbone()，args.masks 的 default='False'
    backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)  # args.backbone 的 default='resnet50'
    model = Joiner(backbone, position_embedding)  # 将 backbone 和 position_embedding 合并一起
    model.num_channels = backbone.num_channels
    return model
